from transformers import AutoModelForCausalLM
import os
import sys
from transformers import AutoConfig
import sentencepiece
import torch
import time

# 添加当前目录到 Python 路径
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)

# 导入必要的模型类
try:
    from my_pangu.modeling_gptpangu import GPTPanguForCausalLM
    from my_pangu.configuration_gptpangu import GPTPanguConfig
    from my_pangu.tokenization_gptpangu import GPTPanguTokenizer
except ImportError:
    raise ImportError("请确保my-pangu目录中包含所需的模型文件")

class ChatBot:
    def __init__(self, model_path="model", device="auto"):
        if not os.path.exists(model_path):
            raise ValueError(f"模型路径 {model_path} 不存在")

        # 自动选择设备
        if device == "auto":
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device
        
        print(f"使用设备: {self.device}")

        # 加载配置和模型
        self.config = GPTPanguConfig.from_pretrained(model_path)
        self.tokenizer = GPTPanguTokenizer(
            model_file=os.path.join(model_path, "vocab.model")
        )
        
        # 优化模型加载
        self.model = GPTPanguForCausalLM.from_pretrained(
            model_path,
            config=self.config,
            trust_remote_code=True,
            local_files_only=True,
            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,  # 使用半精度
            low_cpu_mem_usage=True  # 减少CPU内存使用
        )
        
        # 将模型移动到设备
        self.model = self.model.to(self.device)
        
        # 设置模型为评估模式
        self.model.eval()
        
        # 设置生成参数 - 平衡质量和速度
        self.generation_config = {
            'max_new_tokens': 1500,  # 适中的生成长度
            'min_length': 15,       # 适中的最小长度
            'num_return_sequences': 1,
            'pad_token_id': 9,
            'eos_token_id': 9,
            'do_sample': True,
            'top_p': 0.9,           # 适中的top_p
            'top_k': 60,            # 适中的top_k
            'temperature': 0.8,     # 适中的温度
            'repetition_penalty': 1.1,  # 适中的重复惩罚
            'no_repeat_ngram_size': 3,   # 适中的n-gram检查
            'early_stopping': False,    # 不提前停止
            'use_cache': True,          # 保留KV缓存优化
            'num_beams': 1,             # 使用贪心搜索
        }

    def generate_response(self, input_text):
        # 简单的提示词处理
        if not input_text.endswith("？") and not input_text.endswith("?"):
            input_text = input_text + "？"
        
        # 对输入文本进行编码
        inputs = self.tokenizer(input_text, return_tensors="pt")
        
        # 将输入移动到设备
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # 生成回复 - 优化性能
        with torch.no_grad():  # 禁用梯度计算
            # 使用torch.cuda.amp.autocast()进行混合精度推理（如果使用GPU）
            if self.device == "cuda":
                with torch.cuda.amp.autocast():
                    outputs = self.model.generate(
                        inputs['input_ids'],
                        attention_mask=inputs['attention_mask'],
                        **self.generation_config
                    )
            else:
                outputs = self.model.generate(
                    inputs['input_ids'],
                    attention_mask=inputs['attention_mask'],
                    **self.generation_config
                )
        
        # 解码并返回生成的文本
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # 后处理：移除输入文本，只保留生成的部分
        if input_text in response:
            response = response.replace(input_text, "").strip()
        
        # 清理特殊标记
        response = response.replace("<eot>", "").replace("<pad>", "").strip()
        
        # 进一步清理：移除重复的标点符号
        response = response.replace("？？", "？").replace("。。", "。")
        
        return response
    
    def generate_response_with_timing(self, input_text, show_timing=True):
        """带性能计时的生成方法"""
        if show_timing:
            start_time = time.time()
        
        response = self.generate_response(input_text)
        
        if show_timing:
            end_time = time.time()
            generation_time = end_time - start_time
            print(f"⏱️  生成时间: {generation_time:.2f}秒")
            print(f"📊 生成速度: {len(response)/generation_time:.1f} 字符/秒")
        
        return response

def main():
    # 创建聊天机器人实例
    chatbot = ChatBot()
    
    print("欢迎使用PanGu聊天机器人！输入'退出'结束对话。")
    
    while True:
        # 获取用户输入
        user_input = input("\n用户: ")
        
        # 检查是否退出
        if user_input.lower() in ['退出', 'quit', 'exit']:
            print("\n感谢使用，再见！")
            break
        
        # 生成回复
        try:
            response = chatbot.generate_response(user_input)
            print(f"\nPanGu: {response}")
        except Exception as e:
            print(f"\n生成回复时出错: {str(e)}")

if __name__ == "__main__":
    main()